/*
    2D FDTD simulator
    Copyright (C) 2019 Emilia Blåsten

    This program is free software: you can redistribute it and/or
    modify it under the terms of the GNU Affero General Public License
    as published by the Free Software Foundation, either version 3 of
    the License, or (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public
    License along with this program.  If not, see
    <http://www.gnu.org/licenses/>.
*/
/* abctmz.c: Function to apply a second-order ABC to a TMz grid. */

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include "gridtmz.h"
#include "abctmz.h"

/* Define 3D arrays that store the previous values of the fields. For
 * each one of these arrays the three arguments are as follows:
 *
 *  first argument:  displacement back in time
 *  second argument: spatial displacement from the boundary
 *  third argument:  distance from either the bottom (if EzLeft or
 *                   EzRight) or left (if EzTop or EzBottom) side of
 *                   grid
 */

void abcInit(Abc *ab, Grid *g) {
  double temp1, temp2;

  // calculate ABC coefficients
  temp1 = sqrt( g->ezUpdateHcoeff[0][0] * g->hyUpdateEcoeff[0][0] );
  temp2 = 1.0 / temp1 + 2.0 + temp1;
  ab->coef0 = -(1.0 / temp1 - 2.0 + temp1) / temp2;
  ab->coef1 = -2.0 * (temp1 - 1.0 / temp1) / temp2;
  ab->coef2 = 4.0 * (temp1 + 1.0 / temp1) / temp2;

  return;
}

void abc(Abc *ab, Grid *g) {
  int mm, nn;

  // ABC at left side of grid
  for(nn = 0; nn < g->sizeY; nn++) {
    g->ez[0][nn] = ab->coef0 * (g->ez[2][nn] + ab->ezLeft[1][0][nn])
      + ab->coef1 * ( ab->ezLeft[0][0][nn] + ab->ezLeft[0][2][nn]
                            - g->ez[1][nn] - ab->ezLeft[1][1][nn] )
      + ab->coef2 * ab->ezLeft[0][1][nn]
      - ab->ezLeft[1][2][nn];

      // memorize old fields
    for(mm = 0; mm < 3; mm++) {
      ab->ezLeft[1][mm][nn] = ab->ezLeft[0][mm][nn];
      ab->ezLeft[0][mm][nn] = g->ez[mm][nn];
    }
  }

  // ABC at right side of grid
  for(nn = 0; nn < g->sizeY; nn++) {
    g->ez[g->sizeX-1][nn] =
        ab->coef0 * (g->ez[g->sizeX-3][nn] + ab->ezRight[1][0][nn])
      + ab->coef1 * ( ab->ezRight[0][0][nn] + ab->ezRight[0][2][nn]
                    - g->ez[g->sizeX-2][nn] - ab->ezRight[1][1][nn] )
      + ab->coef2 * ab->ezRight[0][1][nn]
      - ab->ezRight[1][2][nn];

    // memorize old fields
    for(mm = 0; mm < 3; mm++) {
      ab->ezRight[1][mm][nn] = ab->ezRight[0][mm][nn];
      ab->ezRight[0][mm][nn] = g->ez[g->sizeX-1-mm][nn];
    }
  }

  // ABC at bottom of grid
  for(mm = 0; mm < g->sizeX; mm++) {
    g->ez[mm][0] = ab->coef0 * (g->ez[mm][2] + ab->ezBottom[1][0][mm])
      + ab->coef1 * ( ab->ezBottom[0][0][mm] + ab->ezBottom[0][2][mm]
                              - g->ez[mm][1] - ab->ezBottom[1][1][mm] )
      + ab->coef2 * ab->ezBottom[0][1][mm]
      - ab->ezBottom[1][2][mm];

    // memorize old fields
    for(nn = 0; nn < 3; nn++) {
      ab->ezBottom[1][nn][mm] = ab->ezBottom[0][nn][mm];
      ab->ezBottom[0][nn][mm] = g->ez[mm][nn];
    }
  }

  // ABC at top of grid
  for(mm = 0; mm < g->sizeX; mm++) {
    g->ez[mm][g->sizeY-1] =
        ab->coef0 * (g->ez[mm][g->sizeY-3] + ab->ezTop[1][0][mm])
      + ab->coef1 * ( ab->ezTop[0][0][mm] + ab->ezTop[0][2][mm]
                  - g->ez[mm][g->sizeY-2] - ab->ezTop[1][1][mm] )
      + ab->coef2 * ab->ezTop[0][1][mm]
      - ab->ezTop[1][2][mm];

    // memorize old fields
    for(nn = 0; nn < 3; nn++) {
      ab->ezTop[1][nn][mm] = ab->ezTop[0][nn][mm];
      ab->ezTop[0][nn][mm] = g->ez[mm][g->sizeY-1-nn];
    }
  }

return;
}


Abc *abcCreate(int sizeX, int sizeY) {
  // Allocate memory for the Abc struct itself
  Abc *ab;
  ab = (Abc *)calloc(1, sizeof(Abc));
  if(!ab) {
    perror("abcCreate()");
    fprintf(stderr, "Failed to allocate memory for Abc.\n");
    exit(-1);
  }


  // Allocate memory for the large ABC arrays
  // ezAA[T][O][P], T=time, O=orthogonal, P=parallel coord to boundary
  ab->ezLeft = (double ***)calloc(2, sizeof(double**));
  ab->ezLeftTfixed = (double **)calloc(2*3, sizeof(double*));
  ab->ezLeftTOfixed = (double *)calloc(2*3*sizeY, sizeof(double));
  for(int t=0; t<2; t++) {
    ab->ezLeft[t] = ab->ezLeftTfixed + t*3;
    for(int m=0; m<3; m++)
      ab->ezLeftTfixed[t*3+m] = ab->ezLeftTOfixed + (t*3+m)*sizeY;
  }

  ab->ezRight = (double ***)calloc(2, sizeof(double**));
  ab->ezRightTfixed = (double **)calloc(2*3, sizeof(double*));
  ab->ezRightTOfixed = (double *)calloc(2*3*sizeY, sizeof(double));
  for(int t=0; t<2; t++) {
    ab->ezRight[t] = ab->ezRightTfixed + t*3;
    for(int m=0; m<3; m++)
      ab->ezRightTfixed[t*3+m] = ab->ezRightTOfixed + (t*3+m)*sizeY;
  }

  ab->ezTop = (double ***)calloc(2, sizeof(double**));
  ab->ezTopTfixed = (double **)calloc(2*3, sizeof(double*));
  ab->ezTopTOfixed = (double *)calloc(2*3*sizeX, sizeof(double));
  for(int t=0; t<2; t++) {
    ab->ezTop[t] = ab->ezTopTfixed + t*3;
    for(int o=0; o<3; o++)
      ab->ezTop[t][o] = ab->ezTopTOfixed + (t*3+o)*sizeX;
  }

  ab->ezBottom = (double ***)calloc(2, sizeof(double**));
  ab->ezBottomTfixed = (double **)calloc(2*3, sizeof(double*));
  ab->ezBottomTOfixed = (double *)calloc(2*3*sizeX, sizeof(double));
  for(int t=0; t<2; t++) {
    ab->ezBottom[t] = ab->ezBottomTfixed + t*3;
    for(int o=0; o<3; o++)
      ab->ezBottom[t][o] = ab->ezBottomTOfixed + (t*3+o)*sizeX;
  }

  // Check for allocation success
  if(!ab->ezLeft || !ab->ezLeftTfixed || !ab->ezLeftTOfixed ||
     !ab->ezRight || !ab->ezRightTfixed || !ab->ezRightTOfixed ||
     !ab->ezTop || !ab->ezTopTfixed || !ab->ezTopTOfixed ||
     !ab->ezBottom || !ab->ezBottomTfixed || !ab->ezBottomTOfixed) {
    perror("abcInit");
    fprintf(stderr, "Allocation of one of the abc arrays failed\n");
    exit(-1);
  }

  return ab;
}

void abcDestroy(Abc *ab) {
  free(ab->ezLeftTOfixed);
  free(ab->ezLeftTfixed);
  free(ab->ezLeft);

  free(ab->ezRightTOfixed);
  free(ab->ezRightTfixed);
  free(ab->ezRight);

  free(ab->ezTopTOfixed);
  free(ab->ezTopTfixed);
  free(ab->ezTop);

  free(ab->ezBottomTOfixed);
  free(ab->ezBottomTfixed);
  free(ab->ezBottom);
  free(ab);

  return;
}
